import os

from langchain_classic.chains.retrieval_qa.base import RetrievalQA  # 检索问答链
from langchain_community.document_loaders import PyPDFLoader  # PDF文档加载器
from langchain_community.embeddings import FakeEmbeddings  # 备用嵌入模型
from langchain_community.llms import LlamaCpp
from langchain_community.vectorstores import FAISS  # 向量数据库
from langchain_core.prompts import PromptTemplate  # 提示模板
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter  # 文本分割器

os.environ["http_proxy"] = "http://127.0.0.1:7890"
os.environ["https_proxy"] = "http://127.0.0.1:7890"

MODEL_PATH = "F:/MODEL/lmstudio-community/Qwen3-8B-GGUF/Qwen3-8B-Q8_0.gguf"  # 本地模型路径
PDF_PATH = "../example.pdf"  # PDF文档路径

"""
问答系统 极简模式 用于理解问答系统的基本单元
"""


def load_qwen_model():
    print("⏳ 正在加载Qwen模型...")
    return LlamaCpp(
        model_path=MODEL_PATH,  # 模型文件路径
        n_ctx=32768,  # 上下文长度，决定模型能处理的最大token数量。  通过ollama show model 获取
        n_batch=64,  # 批量处理大小
        n_threads=8,  # CPU线程数  分配 8 个 CPU 线程用于模型计算
        n_gpu_layers=-1,  # GPU层设置 --CUDA可用时使用所有GPU层
        temperature=0.3,  # 控制输出随机性 0.3 表示较低的创造性以保证准确性
        max_tokens=1024,  # 最大生成token数
        top_p=0.9,  # nucleus sampling参数  保留累积概率最高的 90% token
        repeat_penalty=1.2,  # 重复惩罚因子 增加重复惩罚以减少重复输出
        verbose=False  # 详细输出
    )


def create_qa_system(llm):
    print("⏳ 创建文档问答系统...")
    loader = PyPDFLoader(PDF_PATH)
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,  # 每个文本块大小
        chunk_overlap=200,  # 文本块重叠大小
        separators=["\n\n", "\n", "。", "！", "？"]  # 中文分段符
    )
    docs = text_splitter.split_documents(loader.load())
    embeddings = load_embeddings()
    vector_db = FAISS.from_documents(docs, embeddings)
    prompt = create_prompt_template()
    return RetrievalQA.from_chain_type(
        llm=llm,  # 大语言模型
        chain_type="stuff",  # 链类型
        retriever=vector_db.as_retriever(search_kwargs={"k": 4}),  # 检索器，检索4个相关片段
        return_source_documents=True,  # 返回源文档
        chain_type_kwargs={"prompt": prompt},  # 使用自定义提示模板
        verbose=False  # 详细输出False即极简输出
    )


def load_embeddings():
    print("🔍 正在加载嵌入模型...")

    model_name = "BAAI/bge-small-zh-v1.5"
    model_kwargs = {"device": "cpu"}
    # 设置离线模式环境变量重本地默认路径加载模型  C:\\Users\Administrator\.cache\huggingface\hub
    os.environ["HF_HUB_OFFLINE"] = "1"
    embeddings = HuggingFaceEmbeddings(
        model_name=model_name,  # 模型名称
        model_kwargs=model_kwargs,  # 模型参数
        encode_kwargs={'normalize_embeddings': True}  # 编码参数
    )
    if embeddings is None:
        print("⚠️  所有嵌入模型加载失败，使用基础的文本嵌入方法...")
        embeddings = FakeEmbeddings(size=128)
    return embeddings


def create_prompt_template():
    template = """
    请根据以下上下文回答问题：
    
    {context}
    
    问题：{question}
    
    要求：
    1. 使用简体中文回答
    2. 答案必须基于上下文
    3. 如果上下文未提及，回答"未找到相关信息"
    4. 保持回答简洁专业
    5. 只输出一次答案，不要重复
    6. 直接给出答案，不要添加额外解释
    7. 答案应尽可能简短，但要完整
    """
    return PromptTemplate(template=template, input_variables=["context", "question"])


def main():
    qwen_llm = load_qwen_model()
    print("✅ Qwen模型加载完成")

    qa_chain = create_qa_system(qwen_llm)
    print("✅ 问答系统创建完成")

    # 示例问答
    questions = [
        "今年的营业收入是多少？",
        "今年的净利润是多少？"
    ]

    # 遍历问题并获取答案
    for question in questions:
        print(f"\n❓ 问题: {question}")
        # 调用问答链获取答案
        result = qa_chain.invoke({"query": question})
        # 清理重复的答案内容
        answer = result['result'].strip()
        # 移除重复的说明内容
        if "### 说明：" in answer:
            answer = answer.split("### 说明：")[0]
        # 移除重复的答案行
        lines = answer.split('\n')
        unique_lines = []
        for line in lines:
            line_clean = line.strip()
            if line_clean and line_clean not in unique_lines:
                unique_lines.append(line_clean)
        # 限制输出行数，避免重复
        final_answer = '\n'.join(unique_lines[:5])  # 最多显示5行
        print(f"💡 答案: {final_answer}")
        # 显示答案来源页码
        if result['source_documents']:
            print(f"📌 来源页码: {result['source_documents'][0].metadata['page'] + 1}")


# 程序入口点
if __name__ == "__main__":
    main()
